#include "SymbolTable.h"
#include "ASTNode.h"

// ==================== SymbolTable 成员函数 ====================

void SymbolTable::enterScope() {
    currentScope++;
}

void SymbolTable::exitScope(bool removeSymbols) {
    if (removeSymbols) {
        // 删除当前作用域的所有符号
        for (auto it = entries.begin(); it != entries.end(); ) {
            if (it->second.scopeLevel == currentScope) {
                it = entries.erase(it);
            } else {
                ++it;
            }
        }
    }
    currentScope--;
}

bool SymbolTable::addVariable(const std::string& name, const std::string& dataType) {
    if (entries.find(name) != entries.end() && entries[name].scopeLevel == currentScope) {
        return false; // 重复定义
    }

    SymbolEntry entry;
    entry.name = name;
    entry.type = SymbolType::VARIABLE;
    entry.dataType = dataType;
    entry.scopeLevel = currentScope;
    entry.stackOffset = stackPointer;

    stackPointer += 4; // 假设 int 占 4 字节
    entries[name] = entry;
    return true;
}

bool SymbolTable::addFunction(const std::string& name, const std::string& returnType) {
    if (entries.find(name) != entries.end()) {
        return false; // 函数已定义
    }

    SymbolEntry entry;
    entry.name = name;
    entry.type = SymbolType::FUNCTION;
    entry.dataType = returnType;
    entry.scopeLevel = 0; // 函数是全局作用域
    entry.stackOffset = 0;

    entries[name] = entry;
    return true;
}

bool SymbolTable::addParameter(const std::string& name, const std::string& dataType) {
    if (entries.find(name) != entries.end() && entries[name].scopeLevel == currentScope) {
        return false; // 重复定义
    }

    SymbolEntry entry;
    entry.name = name;
    entry.type = SymbolType::PARAMETER;
    entry.dataType = dataType;
    entry.scopeLevel = currentScope;
    entry.stackOffset = stackPointer;

    stackPointer += 4; // 参数也占 4 字节
    entries[name] = entry;
    return true;
}

SymbolEntry* SymbolTable::lookup(const std::string& name) {
    auto it = entries.find(name);
    return (it != entries.end()) ? &(it->second) : nullptr;
}

void SymbolTable::print() const {
    for (const auto& entry : entries) {
        std::cout << "Name: " << entry.second.name
                  << ", Type: " << static_cast<int>(entry.second.type)
                  << ", DataType: " << entry.second.dataType
                  << ", Scope: " << entry.second.scopeLevel
                  << ", Offset: " << entry.second.stackOffset
                  << std::endl;
    }
}

int SymbolTable::getStackSize() const {
    return stackPointer;
}

void SymbolTable::resetStackPointer() {
    stackPointer = 0; // 开始新函数时重置栈指针
}

// ==================== SymbolTableBuilder 成员函数 ====================

void SymbolTableBuilder::build(ASTNode* root) {
    if (!root || root->type != "CompUnit") return;

    for (const auto& child : root->children) {
        if (child->type == "FuncDef") {
            processFunction(child.get());
        }
    }
}

void SymbolTableBuilder::printSymbolTable() {
    std::cout << "===== Symbol Table =====" << std::endl;
    globalTable.print();
}

void SymbolTableBuilder::processFunction(ASTNode* funcDef) {
    // 1. 添加函数到全局符号表
    std::string returnType = funcDef->children[0]->type;
    std::string funcName = funcDef->children[1]->value;
    currentTable->addFunction(funcName, returnType);

    // 2. 进入函数作用域（同时包含参数和局部变量）
    currentTable->enterScope();
    currentTable->resetStackPointer(); // 重置栈指针

    // 3. 处理参数（与局部变量在同一作用域）
    if (funcDef->children.size() > 2) {
        ASTNode* paramList = funcDef->children[2].get();
        for (const auto& param : paramList->children) {
            if (param->type == "Param" && param->children.size() > 0) {
                std::string paramName = param->children[0]->value;
                currentTable->addParameter(paramName, "int");
            }
        }
    }

    // 4. 处理函数体（与参数在同一作用域）
    if (funcDef->children.size() > 3) {
        // 直接处理块内的语句，不再创建新的作用域
        for (const auto& stmt : funcDef->children[3]->children) {
            processStatement(stmt.get());
        }
    }

    // 5. 退出函数作用域
    currentTable->exitScope(false);
}

void SymbolTableBuilder::processBlock(ASTNode* block) {
    // 进入新作用域（对于函数体外的块）
    currentTable->enterScope();
    for (const auto& stmt : block->children) {
        processStatement(stmt.get());
    }
    currentTable->exitScope(false); // 保留符号
}
void SymbolTableBuilder::processStatement(ASTNode* stmt) {
    if (stmt->type == "DeclAssignStmt") {
        // 变量声明并赋值
        std::string varName = stmt->children[1]->value;
        currentTable->addVariable(varName, "int");
    } 
    else if (stmt->type == "AssignStmt") {
        // 变量赋值
        std::string varName = stmt->children[0]->value;
        // 检查变量是否已声明
        if (!currentTable->lookup(varName)) {
            std::cerr << "Error: Variable '" << varName << "' not declared" << std::endl;
        }
    }
    else if (stmt->type == "Block") {
        processBlock(stmt);
    }
    else if (stmt->type == "IfStmt" || stmt->type == "IfElseStmt") {
        // 处理条件表达式
        processExpression(stmt->children[0].get());
        // 处理then块
        processStatement(stmt->children[1].get());
        // 处理else块（如果有）
        if (stmt->type == "IfElseStmt") {
            processStatement(stmt->children[2].get());
        }
    }
    else if (stmt->type == "WhileStmt") {
        // 处理条件表达式
        processExpression(stmt->children[0].get());
        // 处理循环体
        processStatement(stmt->children[1].get());
    }
    // 其他语句类型可以类似处理
}

void SymbolTableBuilder::processExpression(ASTNode* expr) {
    // 主要检查变量引用
    if (expr->type == "VarRef") {
        std::string varName = expr->value;
        if (!currentTable->lookup(varName)) {
            std::cerr << "Error: Variable '" << varName << "' not declared" << std::endl;
        }
    }
    // 递归处理子表达式
    for (const auto& child : expr->children) {
        processExpression(child.get());
    }
}